"""
========================================================================
FILE: ensemble_field_function.py

LAST MODIFIED: 23 September 2009

DESCRIPTION:
Classes for combining information and functions of mesh topology
(topology.mesh.mesh_ensemble, element.element_types), basis 
(basis.basis), and mapping (mapper.mapper)
========================================================================
"""
"""
BEGIN LICENSE BLOCK 
Version: MPL 1.1/GPL 2.0/LGPL 2.1

The contents of this file are subject to the Mozilla Public License Version
1.1 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at
http://www.mozilla.org/MPL/

Software distributed under the License is distributed on an "AS IS" basis,
WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
for the specific language governing rights and limitations under the
License.

The Original Code is fieldwork.

The Initial Developer of the Original Code is Ju Zhang.
Portions created by the Initial Developer are Copyright (C) 2009
the Initial Developer. All Rights Reserved.

Contributor(s):

Alternatively, the contents of this file may be used under the terms of
either the GNU General Public License Version 2 or later (the "GPL"), or
the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
in which case the provisions of the GPL or the LGPL are applicable instead
of those above. If you wish to allow use of your version of this file only
under the terms of either the GPL or the LGPL, and not to allow others to
use your version of this file under the terms of the MPL, indicate your
decision by deleting the provisions above and replace them with the notice
and other provisions required by the GPL or the LGPL. If you do not delete
the provisions above, a recipient may use your version of this file under
the terms of any one of the MPL, the GPL or the LGPL.

END LICENSE BLOCK
"""

import scipy
from basis import basis
from topology import mesh
from topology import element_types
import mapper
import copy

class ensemble_field_function:
	""" Class that combines basis, mesh and mapper to create a field.
	Acts as a wrapper for methods in mesh and mapper and handles field
	evaluation.
	"""
	
	def __init__( self, name, dimensions, debug = 0 ):
		self.name = name				# string
		self.dimensions = dimensions	# integer
		self.basis = None				# basis object
		self.basis_type = None			# string
		self.mesh = None				# mesh_ensemble object
		self.mapper = mapper.mapper()	# parameter mapper object
		self.parameters = None
		self.subfields = {}				# {subfield number: subfield}
		self.submesh_map = {}			# {submesh: corresponding subfield}
		self.subfield_counter = 0
		self.is_element = False
		
		self.debug = debug
	
	#==================================================================#	
	def set_basis( self, types ):
		""" Instantiate a basis object of defined type. types is either a
		string or list of strings for multiple basis functions. If a 
		string, assume same basis type for all dimensions.
		"""
		
		if not isinstance( types, list ):
			types = [ types ] * self.dimensions
		
		if len( types ) != self.dimensions:
			print 'ERROR: set_basis: incorrect number of basis types, need', self.dimensions
			return
		else:
			self.basis = basis.make_basis( self.dimensions, types )
			if self.basis:
				self.basis_type = self.basis.type
				return 1
			else:
				return
				
	#==================================================================#		
	def set_new_mesh( self, name ):
		""" Instantiate a new mesh_ensemble object.
		"""
		
		self.mesh = mesh.mesh_ensemble( name, self.dimensions )
		if self.mesh:
			return 1
		else:
			return
	
	#==================================================================#
	def set_mesh( self, mesh ):
		""" Set an existing mesh object to be the field topology
		automatically updates global_map.
		"""
		
		self.mesh = mesh
		self.map_parameters()
		return 1

	#==================================================================#
	def add_element( self, element ):
		""" Add an existing element (element or ensemble field function 
		object) to the mesh. returns new element number for convenience.
		"""
		
		if element.is_element:
			n = self.mesh.add_element( element ) 
			if n == None:
				print 'ERROR: ensemble_field_function.add_element: unable to add element to mesh'
				return
		else:
			n = self._add_subfield( element )
			if n == None:
				print 'ERROR: ensemble_field_function.add_element: unable to add subfield to mesh'
				return
					
		return n
		
	#==================================================================#	
	def create_elements( self, type, number_of_elements ):
		""" Create new elements of the specified type and add to mesh.
		Returns a tuple of new element numbers for convenience.
		"""
		
		n_all = []
		for i in range( number_of_elements ):
			# create element
			element = element_types.create_element( type )
		
			if element:
				n = self.mesh.add_element( element )
				if n == None:
					print 'ERROR: ensemble_field_function.create_element: unable to add element to mesh'
					return
				n_all.append(n)
			else:
				print 'ERROR ensemble_field_function.create_element: unable to create element'
				return
				
		return tuple(n_all)
	
	#==================================================================#
	def _add_subfield( self, subfield ):
		""" Add an existing ensemble field function object as a subfield
		update self.mesh connectivity.
		"""
		
		self.submesh_map[ subfield.mesh ] = subfield
		self.subfields[self.subfield_counter] = subfield
		n = self.mesh.add_element( subfield.mesh )
		
		self.subfield_counter += 1
		return n
		
	#==================================================================#
	def remove_element( self, element_number ):
		""" Remove element element_number from the mesh.
		"""
		
		if self.mesh.remove_element( element_number ):
			# removed element from the mapper
			if self.mapper.remove_element( element_number ):
				return 1
				
		return
			
	#==================================================================#
	def connect_element_points( self, points ):
		""" Wrapper for mesh_ensemble's connect method.
		"""
		
		if self.mesh.connect( points ):
			return 1
		else:
			return
	
	#==================================================================#s
	def connect_to_hanging_point( self, host_element, element_coordinates, connected_points ):
		""" Wrapper for mesh_ensemble's connect_to_hanging_point method.
		"""
		
		if self.mesh.connect_to_hanging_point( host_element, element_coordinates, connected_points ):
			return 1
		else:
			return
		
	#==================================================================#
	def map_parameters( self ):
		""" Use mapper object to generates global node and elements 
		numbers using connectivity information.
		"""
		
		self.mapper.set_parent_field( self )
		if self.mapper.do_mapping():
			return 1
		else:
			print 'ERROR: ensemble_field_function.map_parameters failed'
			return
	
	#==================================================================#
	def get_mapping( self ):
		""" Returns the mapping dictionaries generated by the mapper.
		"""
		
		en_to_el = copy.deepcopy( self.mapper._ensemble_to_element_map )
		el_to_en = copy.deepcopy( self.mapper._element_to_ensemble_map )
		
		return ( en_to_el, el_to_en )
	
	#==================================================================#
	def get_number_of_ensemble_points( self ):
		""" Returns the number of ensemble points in the field by
		querying the mapper.
		"""
		
		return self.mapper.get_number_of_ensemble_points()
			
	#==================================================================#		
	def set_parameters( self, parameters ):
		""" Set the parameters for the field. Not sure if this is
		necessary
		"""
		
		# check parameter length
		if len( parameters ) != self.mapper.get_number_of_ensemble_points():
			print 'ERROR: ensemble_field_function:set_parameters: wrong number of parameter sets, there are', self.mapper.get_number_of_ensemble_points(), 'nodes'
			return
		else:
			self.parameters = parameters
			return 1
	
	#==================================================================#	
	def evaluate_field_at_element_node( self, element, point ):
		# gets the values at element int:element point int:point
		# by querying mapper
		
		element_parameters = self.mapper.get_element_parameters( element, self.parameters )
		
		if element_parameters == None:
			print 'ERROR: ensemble_field_function.evaluate_field_at_element_point: unable to get element parameters from mapper'
			return None
		else:
			try:
				point_values = element_parameters[point]
			except IndexError:
				print 'ERROR: ensemble_field_function.evaluate_field_at_element_point: invalid point in element'
				return None
			else:
				return point_values[0]
	
	#==================================================================#
	def evaluate_field_in_mesh( self, density, parameters = None ):
		""" Evaluates the field over the whole mesh, i.e. in all 
		elements. Returns a list of field values.
		
		density is the sampling density per dimension in each element.
		parameters is a list of field parameters, with a list for each 
		ensemble point. If parameters is not passed, will look to see if
		parameters have been set. If not, returns error.
		"""
		
		if parameters == None:
			parameters = self.parameters
		
		# check that there are parameters
		if parameters == None:
			print 'ERROR: ensemble_field_function.evaluate_field_in_mesh: no parameters passed or set'
			return
		
		# check for right number of density values
		if isinstance( density, int ):
			density = [density] * self.dimensions
		elif len( density ) != self.dimensions:
			print 'ERROR: evaluate_element: needed', self.dimensions, 'density values. Got', len( density )
			return
	
		field_values = []
		element_field_values = None
	
		for element_number in self.mesh.elements.keys():
			# get element
			element = self.mesh.elements[ element_number ] 
			# for each element, get element_parameters
			element_parameters = self.mapper.get_element_parameters( element_number, parameters )
			
			if element.is_element:
				# if element is local, evaluate using evaluate_field_in_element
				element_field_values = self._evaluate_element( element, element_parameters, density )
			else:
				# else element is a sub mesh, evaluate the corresponding subfield
				element_field_values = self.submesh_map[ element ].evaluate_field_in_mesh( density, element_parameters )
			
			# store evaluated values
			if not element_field_values:
				print 'WARNING: ensemble_field_function.evaluate_field_in_mesh: evaluation failed in element', element_number
			else:
				field_values += element_field_values
		
		return field_values	
	
	#==================================================================#
	#~ def evaluate_field_in_element( self, element_number, density ):
		#~ # evaluates the field in an element defined by element_number
		#~ 
		#~ # check for right number of density values
		#~ if isinstance( density, int ):
			#~ density = [density] * self.dimensions
		#~ elif len( density ) != self.dimensions:
			#~ print 'ERROR: evaluate_element: needed', self.dimensions, 'density values. Got', len( density )
			#~ return 0
	#~ 
		
	#==================================================================#	
	def _evaluate_element( self, element, element_parameters, density ):
		""" Evaluates a given element with its parameters.
		
		element is an element object
		element_parameters is a list parameters containing a list for 
		each element node.
		density is a list of number of points evaluated in each element direction
		"""
		
		# check for right number of density values
		if isinstance( density, int ):
			density = [density] * self.dimensions
		elif len( density ) != self.dimensions:
			print 'ERROR: evaluate_element: needed', self.dimensions, 'density values. Got', len( density )
			return None
		
		# initialise
		eval_coords = []
		bases = []
		eval_values = []
		#~ element = self.mesh.elements[element_number]
		#~ element_parameters = self._get_element_parameters( element_number )
		
		if self.debug:
			print 'element_parameters:', element_parameters
		
		# generate evaluation points
		for i in range( self.dimensions ):
			eval_coords.append( scipy.linspace( element.interior[i][0], element.interior[i][1], density[i] ) )
		
		# evaluate basis functions at each evaluation point
		if self.dimensions == 1:
			for i in eval_coords[0]:
				if element.is_interior( [i] ):
					bases.append( self.basis.eval( [i] ) )
		elif self.dimensions == 2:
			for i in eval_coords[0]:
				for j in eval_coords[1]:
					if element.is_interior( [i,j] ):
						bases.append( self.basis.eval( [i,j] ) )
		elif self.dimensions == 3:
			for i in eval_coords[0]:
				for j in eval_coords[1]:
					for k in eval_coords[2]:
						if element.is_interior( [i,j,k] ):
							bases.append( self.basis.eval( [i,j,k] ) )	
		else:
			print 'ERROR: ensemble_field_function.evaluate_element: unsupported dimensions'
			return None
		
		if self.debug:
			print '1st basis', bases[0]
		
		# inner product each set of basis values with element parameters
		for basis in bases:
			eval_values.append( scipy.dot( basis, element_parameters ) )
			
		return eval_values	
	
	#==================================================================#
	def get_element_point_evaluator( self, element_number ):
		""" returns an object for evaluating points within element
		element_number
		"""
		
		# check for valid element_number
		try:
			element = self.mesh.elements[ element_number ]
		except KeyError:
			print 'ERROR ensemble_field_function.get_element_point_evaluator: invalid element'
			return None
		else:
			# if element is an element....
			if element.is_element:
				# get element parameters
				parameters = self.mapper.get_element_parameters( element_number )
				# create evaluation object
				evaluator = element_point_evaluator( element, parameters, self.basis )
			# else element is a subfield
			else:
				#~ evaluator = self.submesh_map[ element ].get_element_point_evaluator( element_number )
				print 'ERROR: ensemble_field_function.get_element_point_evaluator: element', element_number, 'is a subfield'
				return None

			return evaluator

	#==================================================================#
	#~ def evaluate_field_at_element_point( self, element_number, element_coordinates, element_parameters ):
		#~ # element_number is an integer
		#~ # element_coordinates should be a list or tuple of xi values
		#~ # 	e.g. [xi1, xi2]
		#~ # element_parameters should be a list of parameters for each
		#~ # 	element point
		#~ #	e.g. [ [p1 parameters], [p2 parameters]...]
#~ 
		#~ element = self.mesh.elements[ element_number ]
#~ 
#~ 
#~ 
		#~ # check coordinates are in element domain
		#~ 
		#~ if element.is_interior( element_coordinates ):
				#~ 
			#~ u_vector = []
			#~ 
			#~ if ( self.basis_type == 'cubic_hermite' ) and self.dimensions == 2:
				#~ # special mapping of element parameters to the tensor product
				#~ # of basis functions for 2D cubic hermite
				#~ # should be improved!
				#~ # goes [1 2 5 6 3 4 7 8 9 10 13 14 11 12 15 16]
				#~ u_vector = []
				#~ for i in range( 0, self.dimensions+1, 2 ):
					#~ for j in range( 0, self.dimensions+1, 2 ):
						#~ u_vector += list( element_point_values[i:i+self.dimensions, j:j+self.dimensions].ravel() )
				#~ 
			#~ else:
				#~ u_vector = element_point_values.ravel()
			#~ 
			#~ if self.debug:
				#~ print 'element_point_values:', element_point_values
				#~ print 'u_vector:', u_vector
			#~ 
			#~ #
			#~ for i in range( len( element_coordinates ) ):
				#~ # evaluate basis function at element_coordinates
				#~ phi = self.basis.eval( element_coordinates[i] )
				#~ 
				#~ if self.debug:
					#~ print 'phi:', phi
			#~ 
				#~ if not isinstance( phi, int ):
					#~ # inner-product with basis
					#~ field_value = scipy.dot( phi, u_vector )
					#~ if self.debug:
						#~ print 'field_values:', field_values
				#~ else:
					#~ print 'ERROR: basis.eval'
					#~ return 0
					#~ 
				#~ field_values.append( field_value )
				#~ 
				#~ 
			#~ return field_values
	
	#==================================================================#	


class element_point_evaluator:
	""" Class for objects evaluating the field within an element
	"""
	
	def __init__( self, element, parameters, basis_function ):
		self.element = element
		self.parameters = parameters
		self.dimensions = self.element.dimensions
		self.basis = basis_function
		
	def evaluate_at_element_point( element_coordinates ):
		""" Evaluates the objects element at the given element
		coordinates. element_coordinates is a list.
		"""
		
		# evaluate basis function
		basis_coeff = self.basis.eval( element_coordinates )
		if basis_coeff != None:
			try:
				field_value = scipy.dot( basis_coeff, self.parameters )
				return field_value
			except ValueError:
				print 'ERROR: element_point_evaluator.evaluate_at_element_point: dot product failed, basis_coeff and parameters misaligned'
				return None
				
		else:
			print 'ERROR: element_point_evaluator.evaluate_at_element_point: unable to evaluate basis function at element coordinates:', element_coordinates
			return None

			
